"""
python fujian2_IterativeImputer_BayesianRidge_fillBlank_single_category.py
"""

import os
import json
import pandas as pd
import numpy as np
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
from sklearn.linear_model import BayesianRidge
import matplotlib.pyplot as plt

def IterativeImputer_BayesianRidge_fillBlank_single_category(input_json_path, output_json_path):
    # 创建输出目录
    output_folder = os.path.dirname(output_json_path)
    graph_folder = os.path.join(output_folder, 'graph')
    os.makedirs(output_folder, exist_ok=True)
    os.makedirs(graph_folder, exist_ok=True)
    
    # 读取输入数据
    with open(input_json_path, 'r') as f:
        data = json.load(f)
    
    # 转换数据为DataFrame
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)
    
    # 生成完整的日期范围
    full_date_range = pd.date_range(start="2022-07-01", end="2023-06-30")
    
    # 插入空白日期并标记缺失值
    full_df = pd.DataFrame(index=full_date_range)
    full_df = full_df.join(df['sales'])
    
    # 标记缺失值
    full_df['sales'] = full_df['sales'].astype(float)
    
    # 准备数据进行插补
    full_df['is_missing'] = full_df['sales'].isna()
    X = np.arange(len(full_df)).reshape(-1, 1)  # 使用索引作为特征
    y = full_df['sales'].values
    
    # 初始化迭代插补器
    imputer = IterativeImputer(estimator=BayesianRidge(), random_state=0)
    
    # 填补缺失值
    full_df['sales'] = imputer.fit_transform(X, y.reshape(-1, 1)).ravel()
    
    # 提取填补后的数据
    filled_data = full_df.loc["2022-10-01":"2023-03-31"]
    filled_data = filled_data[['sales']].reset_index()
    filled_data.columns = ['date', 'sales']
    filled_data['date'] = filled_data['date'].dt.strftime('%Y/%m/%d')
    
    # 保存填补后的数据到JSON文件
    with open(output_json_path, 'w') as f:
        json.dump(filled_data.to_dict(orient='records'), f, indent=4)
    
    # 绘制图表，显示原始和填补后的数据
    plt.figure(figsize=(12, 6))
    plt.plot(full_df.index, full_df['sales'], label='填补后的数据', color='orange')
    plt.plot(df.index, df['sales'], label='原始数据', color='blue')
    plt.xlabel('日期')
    plt.ylabel('销售量')
    plt.legend()
    plt.title('销售数据填补')
    
    # 保存图表
    graph_path = os.path.join(graph_folder, 'sales_data_fill.png')
    plt.savefig(graph_path)
    plt.close()
    
    print(f"填补后的数据已保存至 {output_json_path}")
    print(f"图表已保存至 {graph_path}")

if __name__ == '__main__':
    input_path = "../fujian/fujian2/groupByCategory/category_category1.json"
    output_path = "../fujian/fujian2/IterativeImputer_BayesianRidge_fix/category_category1_filled.json"
    IterativeImputer_BayesianRidge_fillBlank_single_category(input_path, output_path)
